package com.ibeeking.found.common.config.async;

import com.ibeeking.nematos.constants.enums.LogConstant;
import com.ibeeking.nematos.utils.data.StringUtils;
import org.apache.logging.log4j.ThreadContext;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.util.CollectionUtils;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;

/**
 * @Author ibeeking
 */
public class ThreadContextPoolTaskExecutor extends ThreadPoolTaskExecutor {

    private static final long serialVersionUID = 1L;

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        Map<String, String> context = ThreadContext.getContext();
        return super.submit(() -> {
            T result = null;
            if (!CollectionUtils.isEmpty(context) && StringUtils.isNotBlank(context.get(LogConstant.TRACE_ID))) {
                ThreadContext.putAll(context);
            } else {
                //为空设置新值
                ThreadContext.put(LogConstant.TRACE_ID, UUID.randomUUID().toString().replace("-", ""));
            }
            try {
                result = task.call();
            } finally {
                ThreadContext.clearAll();
            }
            return result;
        });
    }

}
